#!/usr/bin/python

##
## Copyright (C) 2009-2010 Luciano Xumerle <luciano.xumerle@gmail.com>
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.
##

from expression import *
from knowngene import *

import time
import sys

try:
    import pysqlite2.dbapi2 as sqlite
except:
    import sqlite3 as sqlite

# mapfields contains:
# sequence|chromosome|direction|sstart|sstop|mismatches|full_uc_code
def printLocation(kg, mapfields):
   for smatch in mapfields:
       gene=[]
       whats='INTRAGENE'
       if smatch[1].find("chr")!=0:
           gene.append( [kg.getSpliceGeneCached( smatch[1] ), '=', '=', '=', '=', '='] )
           whats='EXON_SPLICE'
       try:
           istart=int(smatch[3])
           istop=int(smatch[4])

           if len(gene)==0:
               gene=kg.getCDSGeneCached(smatch[1], istart, istop)
               if len(gene)>0:
                   whats='CDS'

           if len(gene)==0:
               gene=kg.getUTRGeneCached(smatch[1], istart, istop)
               if len(gene)>0:
                   whats='UTR'

           if len(gene)==0:
               gene=kg.getExonGeneCached(smatch[1], istart, istop)
               if len(gene)>0:
                   whats='EXON'

           if len(gene)==0:
               gene=kg.getGeneGeneCached(smatch[1], istart, istop)
               if len(gene)>0:
                   whats='GENE'

           if len(gene)==0:
               gene=kg.getNearGeneGeneCached(smatch[1], istart, istop)
               if len(gene)>0:
                   whats='NEARGENE'
       except:
           gene=[]

       if len(gene)==0:
           gene=[ ['=', '=', '=', '=', '=', '='] ]

       # uniqs gene matchs
       gene.sort()
       unique=[]

       for gg in gene:
           if len(unique)>0 and gg[0]==unique[-1][0]:
               unique[-1]=[gg[0], '=', '=', '=', '=', '=']
           else:
               unique.append(gg)

       for gg in unique:
           row=[ smatch[0], smatch[1], smatch[2], str(smatch[3]), str(smatch[4]) ]
           for a in gg:
               row.append( str(a) )
           row.append( whats )
           print "\t".join(row)


###########################################################################


def readAlignedFile(kg, vmfFile, onlyMulti, neargene):
   ## start reading vmf file
   totalLine=0
   row=[]
   mapfields=[]

   ## VMF file
   fout = open(vmfFile, "a")
   if ".vmf" in vmfFile:
       for line in open(vmfFile, "r"):
           mapfields = splitEland2line(line, onlyMulti)
           if len(mapfields)<1:
               continue
           printLocation(kg, mapfields)
           totalLine+=1

   ## MAQ file
   elif ".maq" in vmfFile:
       for line in open(vmfFile, "r"):
           mapfields = splitMaqLine(line)
           if len(mapfields)<1:
               continue
           printLocation(kg, [mapfields])
           totalLine+=1

   ## RDS from erange
   elif ".rds" in vmfFile:
       sql=sqlite.connect(vmfFile)
       sql.execute("PRAGMA CACHE_SIZE = 5000000")
       cur = sql.cursor()
       uniqs="""SELECT DISTINCT readID, chrom, sense, start, stop, mismatch, chrom FROM uniqs;"""
       multi="""SELECT DISTINCT readID, chrom, sense, start, stop, mismatch, chrom FROM multi;"""
       splicesL="""SELECT DISTINCT readID, chrom, sense, startL, stopL, mismatch, chrom from splices;"""
       splicesR="""SELECT DISTINCT readID, chrom, sense, startR, stopR, mismatch, chrom from splices;"""
       countList=[uniqs, splicesL, splicesR]
       if onlyMulti:
           countList=[multi]
       for query in countList:
           cur.execute(uniqs);
           for e in cur:
               totalLine+=1
               printLocation(kg, [e])

   ## uniqs count from RDS by countFromRDS.sh script
   ## row example: id|chr10|+|100002092|100002128|
   elif ".count" in vmfFile:
       for line in open(vmfFile, "r"):
           row=line.strip().split("|")
           printLocation(kg, [ [row[0], row[1], row[2], row[3], row[4], row[5], row[1]] ])
           totalLine+=1

   ## BOWTIE
   elif ".bwt" in vmfFile:
       old=''
       for line in open(vmfFile, "r"):
           cur=splitBowtieLine(line)
           if old != cur[0]:
               if (len(mapfields)==1 and not onlyMulti) or (len(mapfields)>1 and onlyMulti):
                   printLocation(kg, mapfields)
                   totalLine+=1
               mapfields=[]
               old=cur[0]
           mapfields.append(cur)
       if (len(mapfields)==1 and not onlyMulti) or (len(mapfields)>1 and onlyMulti):
           printLocation(kg, mapfields)
           totalLine+=1

   ## file field are chr, start, stop, ......
   ## or chr, start, ......
   else:
       for line in open(vmfFile, "r"):
           row=line.strip().split("\t")
           chrom=row[0]
           start=int(row[1])
           stop=0
           try:
               stop=int(row[2])
           except:
               stop=start+1
           mapfields=[ "\t".join(row[3:]), chrom, ".", start, stop, ".", "." ]
           if stop==0 :
               mapfields=[ "\t".join(row[2:]), chrom, ".", start, stop, ".", "." ]
           mapfields=[mapfields]
           printLocation(kg, mapfields)
           totalLine+=1

   kg.printLogMessage(True)
   kg.plog("Line Processed: " + str(totalLine) )
   kg.printLogMessage(False)

###########################################################################

def main():
   """main routine"""

   # pars
   database=''
   vmfFile=''
   neargene=0

   onlyMulti=False

   errorMsg="SYNTAX: " + sys.argv[0] + " [-onlyMulti] <DataBase> <vmf FILE> [nearGene]\n"

   if len(sys.argv)<2:
       print >>sys.stderr, errorMsg
       return
   for a in sys.argv[1:]:
       if a == "-onlyMulti":
           onlyMulti=True
           continue
       if database=='':
           database=a
           continue
       if vmfFile=='':
           vmfFile=a
           continue
       if neargene==0:
           try:
               neargene=int(a)
           except:
               neargene=0

   if database=='' or vmfFile=='':
       print >>sys.stderr, errorMsg
       return

   print >>sys.stderr, time.ctime() + " CACHE DATA"

   kg=knowngene(database)
   kg.createCache()

   # manage neargene value
   if neargene>0:
       kg.createNearGeneCache(neargene)

   print >>sys.stderr, time.ctime() + " START MAP"

   readAlignedFile(kg, vmfFile, onlyMulti, neargene)

   print >>sys.stderr, time.ctime() + " END MAP"

if __name__=='__main__': main()
